from pymongo import MongoClient

from StockDailyFactorGeneralFunctions import StockDailyFactorGeneralFunctions

import datetime
import time

class PeriodFactorGeneralFunctions:
    def __init__(self):
        funcObj = StockDailyFactorGeneralFunctions()
        self.fillType = funcObj.fillType

        self.DBIP = funcObj.DBIP
        self.DBPort = funcObj.DBPort
        self.DBName = funcObj.DBName
        self.colSysData = 'sys_data'

        self.WeeklyDBName = 'WeeklyFactors'
        self.MonthlyDBName = 'MonthlyFactors'
        self.QuarterlyDBName = 'QuarterFactors'

    def RetrieveAllPeriodPointDateList(self, colTransDates, periodType):
        funcObj = StockDailyFactorGeneralFunctions()
        DBIP = funcObj.DBIP
        DBPort = funcObj.DBPort
        DBName = funcObj.DBName
        client = MongoClient(DBIP, DBPort)
        db = client[DBName]
        collection_trans_dates = db[colTransDates]
        dailyDatesRes = collection_trans_dates.find()
        transDatesDict = dailyDatesRes[0]
        transDateStrList = []
        transDatesList = []
        totalTransDateList = []

        for (k, v) in transDatesDict.items():
            if k != '_id':
                totalTransDateList.append(datetime.datetime.strptime(str(k), "%Y%m%d"))
            valueStr = str(v)
            if (valueStr.find(periodType)!=-1):
                transDateStrList.append(str(k))
                transDatesList.append(datetime.datetime.strptime(str(k), "%Y%m%d"))
        return (totalTransDateList, transDatesList)


    def RetrievePeriodPointDateList(self, colTransDates, periodType, strStartDate, strEndDate, offsetDayNum):
        totalTransDateList, allPeriodPointDateList = self.RetrieveAllPeriodPointDateList(colTransDates, periodType)

        try:
            startDate = datetime.datetime.strptime(strStartDate, "%Y%m%d")

            endDate = datetime.datetime.strptime(strEndDate, "%Y%m%d")
        except ValueError:
            print("Wrong date input!")
            return []
        ndx = len(allPeriodPointDateList)

        periodDateList = []
        for i in range(0, ndx):
            if startDate <= allPeriodPointDateList[i] and endDate >= allPeriodPointDateList[i]:
                allDatesNdx = totalTransDateList.index(allPeriodPointDateList[i])
                periodDateList.append(totalTransDateList[allDatesNdx+offsetDayNum])
        return periodDateList


    def RetrievePeriodDateGroupList(self, colTransDates, periodType, strPointDate, offsetDayNum):

        datePointDate = datetime.datetime.strptime(strPointDate, "%Y%m%d")

        totalTransDateList, allPeriodPointDateList = self.RetrieveAllPeriodPointDateList(colTransDates, periodType)

        ndxPointDate = totalTransDateList.index(datePointDate)

        periodDateNdxIntotalTransDateList = []
        for periodDateItem in allPeriodPointDateList:
            oneNdx = totalTransDateList.index(periodDateItem)
            periodDateNdxIntotalTransDateList.append(oneNdx)

        ndx = periodDateNdxIntotalTransDateList.index(ndxPointDate)

        selectedDateStrList = []
        if (ndx>0) :
            startDate = totalTransDateList[periodDateNdxIntotalTransDateList[ndx -1] + 1 + offsetDayNum]
            endDate = totalTransDateList[periodDateNdxIntotalTransDateList[ndx] + offsetDayNum]
            ndxOfStartDate = totalTransDateList.index(startDate)
            ndxOfEndDate = totalTransDateList.index(endDate)
            for i in range(ndxOfStartDate, ndxOfEndDate+1):
                selectedDateStrList.append(totalTransDateList[i].strftime('%Y%m%d'))
        else:
            checkNdx = periodDateNdxIntotalTransDateList[ndx] + offsetDayNum
            if (checkNdx >0):
                startDate = totalTransDateList[0]
                endDate = totalTransDateList[periodDateNdxIntotalTransDateList[ndx] + offsetDayNum]
            else:
                return selectedDateStrList
            ndxOfStartDate = totalTransDateList.index(startDate)
            ndxOfEndDate = totalTransDateList.index(endDate)
            for i in range(ndxOfStartDate, ndxOfEndDate+1):
                selectedDateStrList.append(totalTransDateList[i].strftime('%Y%m%d'))

        return selectedDateStrList

    def updateDataUpdateDate(self, strDBName, strUpdatedKeyName, strUpdatedKeyAlias, currentDateStr = time.strftime('%Y-%m-%d')):
        DBIP = self.DBIP
        DBPort = self.DBPort
        DBName = strDBName
        colSysData = self.colSysData

        client = MongoClient(DBIP, DBPort)
        db = client[DBName]
        collection_sys_data = db[colSysData]
        try:
            collection_sys_data.update({}, {'$set': {strUpdatedKeyName: currentDateStr}})
        except:
            print('Failed to update  date to MongoDB, please check.......')
            return()
        print('Successfully updated' + strUpdatedKeyAlias + 'into MongoDB.')

    def retrieveDataUpdateDate(self, strDBName, strUpdateTimeKey):
        DBIP = self.DBIP
        DBPort = self.DBPort
        DBName = strDBName
        colSysData = self.colSysData

        client = MongoClient(DBIP, DBPort)
        db = client[DBName]
        collection_sys_data = db[colSysData]


        sysDataRes = collection_sys_data.find({}, {strUpdateTimeKey: 1})

        dateListDict = sysDataRes[0]

        dataUpdateDate = dateListDict[strUpdateTimeKey]

        return dataUpdateDate

    def ObtainOutputDBName(self, periodTypeStr):
        FactorDBName = ''
        if periodTypeStr == 'W':
            FactorDBName = self.WeeklyDBName
        elif periodTypeStr == 'M':
            FactorDBName = self.MonthlyDBName
        elif periodTypeStr == 'Q':
            FactorDBName = self.QuarterlyDBName
        return FactorDBName

#obj = PeriodFactorGeneralFunctions()


#print(obj.RetrievePeriodPointDateList('trans_dates', 'M', '20160101', '20170701', 0))

#print(obj.RetrievePeriodDateGroupList('trans_dates', 'W', '20120106', 0))